"""
chroma数据库的使用

"""
import chromadb
# collection 名称（类似数据库里的表名，用于逻辑分组）
collection_name = "test_01"


def init_db_client():
    """
    初始化数据库客户端
    支持两种方式：
    1. PersistentClient：本地持久化存储，数据存到本地文件夹（如 ./testdb）
    2. HttpClient：连接运行中的 Chroma Server（服务端模式）
    """
    # 本地存储方式（单机测试用）
    chroma_client = chromadb.PersistentClient(path="./testdb")

    # 服务方式（需要先运行 `chroma run` 启动服务）
    # chroma_client = chromadb.HttpClient(host='localhost', port=8000)
    return chroma_client


def create_collection(collection_name):
    """
    创建（或获取已存在的）Collection
    Collection 类似于数据库的表，用于存放一类文档
    """
    chroma_client = init_db_client()
    collection = chroma_client.get_or_create_collection(name=collection_name)
    return collection


def add_documents(collection, documents):
    """
    往指定 Collection 中写入数据
    documents: 文档原文列表
    ids: 每条文档的唯一 ID，这里简单用 id0, id1, id2 形式
    """
    collection.add(
        documents=documents,  # 文本数据（原始文档内容）
        ids=[f"id{i}" for i in range(len(documents))]  # 生成 ID
    )


def db_test():
    """测试数据库写入与查询"""
    # 获取 / 创建一个 collection
    collection = create_collection(collection_name)

    # 要插入的样例数据（关于马云的基本信息）
    datas = [
        "基本信息",
        "姓名：马云（英文名：Jack Ma）",
        "出生日期：1964年9月10日",
        "出生地：中国浙江省杭州市",
        "业院校：杭州师范大学（原杭州师范学院），英语专业",
        "花名：风清扬（源自金庸小说《笑傲江湖》中的神秘高手风清扬）"
    ]
    # 写入数据库
    add_documents(collection, datas)

    # 查询：问“马云的花名是什么”
    # 语义相似度的“距离值”，越小表示越相关。
    res = collection.query(
        query_texts=["马云的花名是什么"],  # 查询问题
        n_results=5  # 返回最相关的前 5 条结果
    )
    print(res)


# 程序入口
db_test()